from FlagEmbedding import BGEM3FlagModel, FlagReranker
import torch

def change_(x):
    t = list()
    t.append(x)
    return t

def change_2(x):
    if '？' in x:
        t = [i+'？' for i in x.split('？') if x != '']
    elif '?' in x:
        t = [i+'?' for i in x.split('?') if x != '']
    elif '。' in x:
        t = [i+'。' for i in x.split('。') if x != '']
    elif '.' in x:
        t = [i+'.' for i in x.split('.') if x != '']
    else:
        t = x
    return t

def combine_(x1,x2):
    return[[a,b] for a in x1 for b in x2]
class Embedding():
    def __init__(self):
        model_path1="/nfs1/56dev/root/dxj/AI问答/bge-m3"
        model_path2="/nfs1/56dev/root/.cache/huggingface/hub/models--BAAI--bge-reranker-v2-m3/snapshots/12e974610ba9083ed95f3edf08d7e899581f4de4"
        self.bge_model = BGEM3FlagModel(model_path1, use_fp16=False)
        self.reranker = FlagReranker(model_path2, use_fp16=True)

    def input_embed1(self, text):
        """文本向量化"""
        embeddings_1 = self.bge_model.encode(text, batch_size=12, max_length=128)['dense_vecs']
        return embeddings_1
    def input_embed(self, text):
        """文本向量化"""
        embeddings_1 = self.bge_model.encode(text, return_dense=True, return_sparse=True, return_colbert_vecs=True)
        return embeddings_1['colbert_vecs']
    def get_score(self,embeddings1,embeddings2):
        return self.bge_model.colbert_score(embeddings1,embeddings2)

    def get_topK(self, pairs:list, k=3):
        """返回句子对重排序后的前k个句子索引"""
        scores = self.reranker.compute_score(pairs, normalize=True)
        score_dict = dict()
        for i in range(len(scores)):
            score_dict[i] = scores[i]
        sorted_indexes = sorted(score_dict.items(), key=lambda item: item[1], reverse=True)
        sorted_indexes = [index for index, _ in sorted_indexes] # 提取排好序的索引
        return sorted_indexes[:k],max(scores)
    def main(self,input_):

        recommendedDialogue=input_["recommendedDialogue"]
        learnerDialogue=input_["learnerDialogue"]
        print(self.input_embed(recommendedDialogue))
        print(self.input_embed(learnerDialogue))


        dense_score = self.input_embed1(recommendedDialogue) @ self.input_embed1(learnerDialogue).T
        colbert_score = self.get_score(self.input_embed(recommendedDialogue),self.input_embed(learnerDialogue))
        colbert_score = colbert_score.item()
        rerank_score = self.get_topK(pairs = combine_(change_(recommendedDialogue),change_2(learnerDialogue)))[1]
        print(dense_score,colbert_score,rerank_score)

        # data = 0
        # for i in [dense_score,colbert_score,rerank_score]:
        #     if i >= data:
        #         data = i
        # output_json = dict()
        # # output_json['data'] = data
        # print(output_json)
        data = float(max(dense_score, colbert_score, rerank_score))

        output_json = {'data': data}
        print(output_json)
        return  output_json

if __name__ == '__main__':
    Embedding()


